"""
MongoDB 事务
"""

from functools import wraps
from types import TracebackType

from pymongo import WriteConcern
from pymongo.read_concern import ReadConcern
from pymongo.synchronous.client_session import ClientSession

from components.global_param import (
    get_ctx,
    update_request_context,
)
from components.logger import log
from my_typing import *


class MongoDBTransaction:
    """ MongoDB | 事务 | 上下文管理器 """

    def __init__(
        self,
        read_concern: str = "majority",  # 默认使用 majority 确保一致性
        write_concern: int | StrOrNone = "majority",  # 默认使用 majority 确保数据安全
    ) -> None:
        """
        注意：MongoDB 事务中，如果某个 PO 对象被多次更新，在事务结束之前，从 DB 读取到的数据永远不变（特别是 操作历史 数据）

        Args:
            read_concern: 读关注
            write_concern: 写关注

        读关注点（ReadConcern）：
            1. local - 返回节点当前数据
               - 最快但可能不是最新
               - 适合非关键查询
            2. available - 返回分片当前数据
               - 快速但可能不一致
               - 适合统计等非精确查询
            3. majority - 返回大多数节点确认的数据
               - 保证一致性
               - 适合需要强一致性的查询
            4. linearizable - 返回线性一致的数据
               - 最强一致性保证
               - 性能开销最大

        写关注点（WriteConcern）：
            1. w: 0 - 不等待写操作确认
               - 最快但最不安全
               - 适合日志等非关键数据
            2. w: 1 （默认） - 等待主节点确认
               - 中等速度和安全性
               - 适合一般业务操作
            3. w: majority - 等待大多数节点确认
               - 最安全但最慢
               - 适合金融交易等关键操作
            4. w: all - 等待所有节点确认
               - 最高安全性
               - 任一节点故障都会导致写入失败

        选择建议：
            - 普通业务：write_concern="1", read_concern="local"
            - 重要业务：write_concern="majority", read_concern="majority"
            - 金融业务：write_concern="all", read_concern="linearizable"
        """

        self.read_concern = read_concern
        self.write_concern = write_concern

    def __enter__(self) -> ClientSession:
        from db.mongodb_db.client import mongo_client

        log.success("[MongoDB][事务]开始")
        self.session: ClientSession = mongo_client.start_session()
        self.session.start_transaction(
            write_concern=WriteConcern(self.write_concern),
            read_concern=ReadConcern(self.read_concern)
        )

        self._inject_session()

        return self.session

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        try:
            if exc_type is None:
                try:
                    self.session.commit_transaction()
                    log.success("[MongoDB][事务]提交成功")

                except Exception as e:
                    self.session.abort_transaction()
                    log.error(f"[MongoDB][事务]提交失败 开始回滚: {str(e)}")
                    raise

            else:
                self.session.abort_transaction()
                log.error("[MongoDB][事务]发生异常 开始回滚")

        finally:
            self.session.end_session()
            log.success("[MongoDB][事务]结束")
            self._clean_session()

    def _inject_session(self) -> None:
        """ 注入 Session """

        log.success("[MongoDB][事务][上下文]注入全局 SESSION")
        update_request_context({"mongo_session": self.session})

    @staticmethod
    def _clean_session() -> None:
        """ 清空 Session """

        log.success("[MongoDB][事务][上下文]清空全局 SESSION")
        update_request_context({"mongo_session": None})


def mongodb_transaction(func):
    """ MongoDB 事务装饰器 """

    @wraps(func)
    def wrapper(*args, **kwargs):
        with MongoDBTransaction():
            result = func(*args, **kwargs)

        return result

    return wrapper


def inject_transaction_session(func):
    """ 向 kwargs 中注入 MongoDB 事务 Session """

    @wraps(func)
    def wrapper(*args, **kwargs):
        if not kwargs.get("session"):
            kwargs["session"] = get_ctx().mongo_session

        return func(*args, **kwargs)

    return wrapper


if __name__ == "__main__":
    from db.mongodb_db.client import mongo_db

    with MongoDBTransaction():  # 无需显式传递 session
        mongo_db.user.insert_one({"name": "Alen"})
        mongo_db.user.insert_one({"name": "Ben"})
